from typing import List, Tuple, Dict, Any
from rainforeLearn.gomoku.v2.configs.config import GomokuDQNConfig
from rainforeLearn.gomoku.v2.train.reward.reward_calculator import RewardCalculator
from rainforeLearn.gomoku.v2.train.statistics.tracker.batch_tracker import BatchTracker


class RewardSystem:
    """奖励系统管理器 - 负责批次奖励和惩罚的计算"""
    
    def __init__(self, config: GomokuDQNConfig):
        self.config = config
        self.calculator = RewardCalculator(config)
        self.batch_stats = BatchTracker()
        
        # 奖励惩罚历史
        self.win_rate_bonus_awarded: List[Tuple[int, float]] = []
        self.batch_penalty_history: List[Tuple[int, float]] = []
        self.penalty_history: List[Tuple[int, float]] = []
        
        # 累计统计
        self.total_batch_bonus = 0.0
        self.total_batch_penalties = 0.0
    
    def add_batch_game_result(self, winner: int) -> None:
        """添加游戏结果到当前批次"""
        self.batch_stats.add_batch_game_result(winner)
    
    def process_batch_completion(self, current_episode: int) -> Tuple[float, float, float]:
        """处理批次完成，返回(胜率奖励, 批次惩罚, 净奖励)"""
        current_batch_win_rate = self.batch_stats.complete_batch()
        
        if current_batch_win_rate == 0.0:
            # 即使没有奖励惩罚，也要记录当前批次的系数
            self.calculator.record_coefficient_for_batch(current_episode)
            return 0.0, 0.0, 0.0
        
        # 先记录当前批次的系数
        self.calculator.record_coefficient_for_batch(current_episode)

        # 计算奖励和惩罚
        win_rate_bonus = self.calculator.calculate_win_rate_bonus(current_batch_win_rate)
        batch_penalty = self._calculate_current_batch_penalty()
        net_reward = win_rate_bonus - batch_penalty
        
        # 动态调整惩罚系数
        self._update_penalty_coefficient(current_episode)

        # 更新统计
        self._update_reward_statistics(current_episode, win_rate_bonus, batch_penalty)
        
        return win_rate_bonus, batch_penalty, net_reward
    
    def _update_penalty_coefficient(self, episode: int) -> None:
        """更新惩罚系数"""
        # 每1个批次评估一次效果并调整系数
        if len(self.batch_stats.win_rates) >= 3 and episode % (self.config.training_stats_batch_freq) == 0:
            effectiveness = self.calculator.calculate_penalty_effectiveness(
                self.batch_stats.win_rates, self.batch_penalty_history
            )
            self.calculator.adjust_penalty_coefficient(effectiveness, episode)

    def _calculate_current_batch_penalty(self) -> float:
        """计算当前批次惩罚"""
        win_rates = self.batch_stats.win_rates
        if len(win_rates) < 2:
            return 0.0
        
        current_rate = win_rates[-1]
        previous_rate = win_rates[-2]
        
        return self.calculator.calculate_batch_penalty(current_rate, previous_rate)
    
    def _update_reward_statistics(self, episode: int, win_rate_bonus: float, batch_penalty: float) -> None:
        """更新奖励统计信息"""
        # 更新累计统计
        self.total_batch_bonus += win_rate_bonus
        self.total_batch_penalties += batch_penalty
        
        # 记录历史
        self.win_rate_bonus_awarded.append((episode, win_rate_bonus))
        self.batch_penalty_history.append((episode, batch_penalty))
        self.penalty_history.append((episode, batch_penalty))
    
    def get_summary(self) -> Dict[str, Any]:
            """获取奖励系统摘要"""
            coefficient_summary = self.calculator.get_coefficient_adjustment_summary()
            
            return {
                'system_status': 'active' if self.batch_stats.win_rates else 'inactive',
                'current_coefficient': self.config.batch_penalty_coefficient,
                'total_batch_bonus': self.total_batch_bonus,
                'total_batch_penalties': self.total_batch_penalties,
                'net_bonus': self.total_batch_bonus - self.total_batch_penalties,
                'recent_batch_win_rates': self.batch_stats.win_rates,
                'coefficient_history': coefficient_summary['coefficient_history'],
                'total_adjustments': coefficient_summary['total_adjustments']
            }